#include "EBFV1_elliptic.h"
#include <time.h>

namespace PRS           // PRS: Petroleum Reservoir Simulator
{
	EBFV1_elliptic::EBFV1_elliptic(){
	}

	EBFV1_elliptic::EBFV1_elliptic(PhysicPropData *ppd,SimulatorParameters *sp,MeshData *md){
		pPPData = ppd;
		pSimPar = sp;
		pMData = md;
		allocateMatVec = false;
	}

	EBFV1_elliptic::~EBFV1_elliptic(){
	}

	// solves system of equation for pressure field
	double EBFV1_elliptic::solver(TMesh* pTMesh_data){
		double startt = MPI_Wtime();
		PetscPrintf(PETSC_COMM_WORLD,"Elliptic solver...\t");
		double cpu_time = .0;
		cpu_time += assemblyMatrices(pTMesh_data);					// assembly matrices
		//cout << "assembly time: " << cpu_time << endl;
		cpu_time += assemblyRHS(pTMesh_data);						// assembly RHS vector (dirichlet and Neumann boundary conditions)
		//cout << "assembly time: " << cpu_time << endl;
		cpu_time += solver_MatFree();								// solver: Ax=y
		//cout << "solver time: " << cpu_time << endl;
		cpu_time += updatePressure(pTMesh_data);					// transfer data solution for data base
		//cout << "updatePressure time: " << cpu_time << endl;
		cpu_time += pressureGradient(pTMesh_data);					// calculate pressure gradient
		//cout << "pressureGradient time: " << cpu_time << endl;
		zeroMatricesAndVectors();									// for the next time step
		PetscPrintf(PETSC_COMM_WORLD,"done.\n");
		double endt = MPI_Wtime();
		return endt-startt;
	}

	double EBFV1_elliptic::solver_MatFree(){
		double startt = MPI_Wtime();
		static bool key = true;
		PetscInt m, n, its;
		ierr = VecGetOwnershipRange(pMatVec->solution,&m,&n); CHKERRQ(ierr);
		PetscInt nLRows = n - m;
		PetscInt numGF = pMData->getNum_GF_Nodes();
		if (key){
			ierr = MatCreateShell(PETSC_COMM_WORLD,nLRows,nLRows,numGF,numGF,pMatVec,&matrix);CHKERRQ(ierr);
			ierr = MatSetFromOptions(matrix); CHKERRQ(ierr);
			ierr = MatShellSetOperation(matrix, MATOP_MULT,(void(*)(void))&EBFV1_elliptic::MatMultUser); CHKERRQ(ierr);
		}
//		ierr = KSP_solver(matrix,pMatVec->G_freeNodes,pMatVec->RHS,pMatVec->solution,pSimPar,PETSC_FALSE,KSPGMRES,PCASM,its); CHKERRQ(ierr);
		KSP_solver(matrix,pMatVec->G_freeNodes,pMatVec->RHS,pMatVec->solution,pSimPar,PETSC_FALSE,KSPGMRES,PCASM,its);
		key = false;
		//ierr = MatDestroy(matrix);CHKERRQ(ierr);
		//printVectorToFile(pMatVec->solution,"SOLUTION.txt");
		double endt = MPI_Wtime();
		return endt-startt;
	}

	double EBFV1_elliptic::updatePressure(TMesh* pTMesh_data){
		double startt = MPI_Wtime();

		PetscScalar *sol, val;
		PetscInt i,m,n,row,col=0;
		PetscInt numGN = pMData->getNum_GNodes();
		Mat mSol,matLocalSolution;

		// create a column matrix to receive pMatVec->solution vector values (mSol)
		ierr = MatCreateMPIAIJ(PETSC_COMM_WORLD,PETSC_DECIDE,PETSC_DECIDE,numGN,1,0,PETSC_NULL,0,PETSC_NULL,&mSol);CHKERRQ(ierr);

		int nLIDs, *IDs_ptr;
		pMData->getRemoteIDs(nLIDs,&IDs_ptr);

		// transference process: from vector to column matrix
		ierr = VecGetArray(pMatVec->solution,&sol);CHKERRQ(ierr);
		ierr = MatSetValues(mSol,pMatVec->nrows,pMatVec->rows,1,&col,sol,INSERT_VALUES);CHKERRQ(ierr);
		ierr = VecRestoreArray(pMatVec->solution,&sol);CHKERRQ(ierr);
		ierr = MatAssemblyBegin(mSol,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
		ierr = MatAssemblyEnd(mSol,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);

		// remote values cannot be gotten from remote matrix positions.
		// transfer (via MatGetSubMatrixRaw) necessary remote values to each process
		// to a second column matrix (matLocalSolution).
		ierr = MatGetSubMatrixRaw(mSol,nLIDs,IDs_ptr,1,&col,PETSC_DECIDE,MAT_INITIAL_MATRIX,&matLocalSolution);CHKERRQ(ierr);
		ierr = MatDestroy(mSol);CHKERRQ(ierr);
		ierr = MatGetOwnershipRange(matLocalSolution,&m,&n);CHKERRQ(ierr);

		// map node ID to its index (For multidomain meshes, ID over sub-domains will have two or more indices assigned. THat's why we use std::multimap)
		row = m;
		int nrows, ncols;
		for (int dom=0; dom<pMatVec->ndom; dom++){
			this->freeNodesIndices.getSize(nrows,ncols);
			for (int j=0; j<nrows; j++){
				int idx = freeNodesIndices.getValue(j,dom);
				VertexData* vdata = pTMesh_data->vertices_list[dom].getValue(idx,0);
				ierr = MatGetValues(matLocalSolution,1,&row,1,&col,&vdata->p);CHKERRQ(ierr);
				row++;
				//cout << row << endl;
			}
		}
//		cout << __LINE__<< endl;
		ierr = MatDestroy(matLocalSolution);CHKERRQ(ierr);
		ierr = MatDestroy(pMatVec->EF);CHKERRQ(ierr);
//		ierr = MatDestroy(pMatVec->E_freeNodes[0]); CHKERRQ(ierr);
//		ierr = MatDestroy(pMatVec->F_freeNodes[0]); CHKERRQ(ierr);
		return MPI_Wtime()-startt;
	}

	int EBFV1_elliptic::zeroMatricesAndVectors(){
		for(int i=0; i<pMatVec->ndom; i++){
			ierr = MatZeroEntries(pMatVec->E[i]); CHKERRQ(ierr);
//			ierr = MatZeroEntries(pMatVec->F[i]); CHKERRQ(ierr);
//			ierr = MatZeroEntries(pMatVec->E_freeNodes[i]); CHKERRQ(ierr);
//			ierr = MatZeroEntries(pMatVec->F_freeNodes[i]); CHKERRQ(ierr);
		}
		ierr = MatZeroEntries(pMatVec->G); CHKERRQ(ierr);
		ierr = MatZeroEntries(pMatVec->G_freeNodes); CHKERRQ(ierr);
		ierr = VecZeroEntries(pMatVec->RHS); CHKERRQ(ierr);
		ierr = VecZeroEntries(pMatVec->solution); CHKERRQ(ierr);
		ierr = VecZeroEntries(pMatVec->z); CHKERRQ(ierr);
	}
}
